# -*- coding: utf-8 -*-
"""
文件存储模块
"""

import json
import csv
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, List, Optional
import pandas as pd
from utils.logger import get_logger, log_storage_info, log_crawler_error
from config.settings import NOVELS_DIR, CHAPTERS_DIR

class FileStorage:
    """文件存储类"""
    
    def __init__(self, base_dir: str = None):
        self.logger = get_logger("FileStorage")
        self.base_dir = Path(base_dir) if base_dir else Path.cwd() / "data"
        self.novels_dir = self.base_dir / "novels"
        self.chapters_dir = self.base_dir / "chapters"
        
        # 创建必要的目录
        self.novels_dir.mkdir(parents=True, exist_ok=True)
        self.chapters_dir.mkdir(parents=True, exist_ok=True)
    
    def save_novel_json(self, novel: Dict[str, Any], filename: str = None) -> bool:
        """
        保存小说数据为JSON格式
        
        Args:
            novel: 小说数据
            filename: 文件名（可选）
            
        Returns:
            是否保存成功
        """
        try:
            if not filename:
                novel_id = novel.get('id', novel.get('novel_id', 'unknown'))
                filename = f"novel_{novel_id}.json"
            
            filepath = self.novels_dir / filename
            
            # 添加保存时间
            novel['saved_at'] = datetime.now().isoformat()
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(novel, f, ensure_ascii=False, indent=2)
            
            log_storage_info(f"小说数据保存成功: {filepath}")
            return True
            
        except Exception as e:
            log_crawler_error(f"保存小说JSON失败: {str(e)}")
            return False
    
    def save_novels_json(self, novels: List[Dict[str, Any]], filename: str = None) -> bool:
        """
        保存小说列表为JSON格式
        
        Args:
            novels: 小说列表
            filename: 文件名（可选）
            
        Returns:
            是否保存成功
        """
        try:
            if not filename:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"novels_{timestamp}.json"
            
            filepath = self.novels_dir / filename
            
            # 添加保存时间
            data = {
                'novels': novels,
                'count': len(novels),
                'saved_at': datetime.now().isoformat()
            }
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            
            log_storage_info(f"小说列表保存成功: {filepath} - 共 {len(novels)} 本")
            return True
            
        except Exception as e:
            log_crawler_error(f"保存小说列表JSON失败: {str(e)}")
            return False
    
    def save_novels_csv(self, novels: List[Dict[str, Any]], filename: str = None) -> bool:
        """
        保存小说列表为CSV格式
        
        Args:
            novels: 小说列表
            filename: 文件名（可选）
            
        Returns:
            是否保存成功
        """
        try:
            if not filename:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"novels_{timestamp}.csv"
            
            filepath = self.novels_dir / filename
            
            if not novels:
                log_crawler_error("小说列表为空，无法保存CSV")
                return False
            
            # 获取所有字段
            fields = set()
            for novel in novels:
                fields.update(novel.keys())
            
            fields = sorted(list(fields))
            
            with open(filepath, 'w', newline='', encoding='utf-8-sig') as f:
                writer = csv.DictWriter(f, fieldnames=fields)
                writer.writeheader()
                writer.writerows(novels)
            
            log_storage_info(f"小说列表CSV保存成功: {filepath} - 共 {len(novels)} 本")
            return True
            
        except Exception as e:
            log_crawler_error(f"保存小说列表CSV失败: {str(e)}")
            return False
    
    def save_novels_excel(self, novels: List[Dict[str, Any]], filename: str = None) -> bool:
        """
        保存小说列表为Excel格式
        
        Args:
            novels: 小说列表
            filename: 文件名（可选）
            
        Returns:
            是否保存成功
        """
        try:
            if not filename:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"novels_{timestamp}.xlsx"
            
            filepath = self.novels_dir / filename
            
            if not novels:
                log_crawler_error("小说列表为空，无法保存Excel")
                return False
            
            # 转换为DataFrame
            df = pd.DataFrame(novels)
            
            # 保存为Excel
            df.to_excel(filepath, index=False, engine='openpyxl')
            
            log_storage_info(f"小说列表Excel保存成功: {filepath} - 共 {len(novels)} 本")
            return True
            
        except Exception as e:
            log_crawler_error(f"保存小说列表Excel失败: {str(e)}")
            return False
    
    def save_chapter_json(self, chapter: Dict[str, Any], novel_id: str, filename: str = None) -> bool:
        """
        保存章节数据为JSON格式
        
        Args:
            chapter: 章节数据
            novel_id: 小说ID
            filename: 文件名（可选）
            
        Returns:
            是否保存成功
        """
        try:
            if not filename:
                chapter_id = chapter.get('id', chapter.get('chapter_id', 'unknown'))
                filename = f"chapter_{novel_id}_{chapter_id}.json"
            
            # 创建小说专属目录
            novel_chapters_dir = self.chapters_dir / novel_id
            novel_chapters_dir.mkdir(parents=True, exist_ok=True)
            
            filepath = novel_chapters_dir / filename
            
            # 添加保存时间
            chapter['saved_at'] = datetime.now().isoformat()
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(chapter, f, ensure_ascii=False, indent=2)
            
            log_storage_info(f"章节数据保存成功: {filepath}")
            return True
            
        except Exception as e:
            log_crawler_error(f"保存章节JSON失败: {str(e)}")
            return False
    
    def save_chapters_json(self, chapters: List[Dict[str, Any]], novel_id: str, filename: str = None) -> bool:
        """
        保存章节列表为JSON格式
        
        Args:
            chapters: 章节列表
            novel_id: 小说ID
            filename: 文件名（可选）
            
        Returns:
            是否保存成功
        """
        try:
            if not filename:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"chapters_{novel_id}_{timestamp}.json"
            
            # 创建小说专属目录
            novel_chapters_dir = self.chapters_dir / novel_id
            novel_chapters_dir.mkdir(parents=True, exist_ok=True)
            
            filepath = novel_chapters_dir / filename
            
            # 添加保存时间
            data = {
                'novel_id': novel_id,
                'chapters': chapters,
                'count': len(chapters),
                'saved_at': datetime.now().isoformat()
            }
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            
            log_storage_info(f"章节列表保存成功: {filepath} - 共 {len(chapters)} 章")
            return True
            
        except Exception as e:
            log_crawler_error(f"保存章节列表JSON失败: {str(e)}")
            return False
    
    def save_chapter_text(self, chapter: Dict[str, Any], novel_id: str, filename: str = None) -> bool:
        """
        保存章节内容为纯文本格式
        
        Args:
            chapter: 章节数据
            novel_id: 小说ID
            filename: 文件名（可选）
            
        Returns:
            是否保存成功
        """
        try:
            if not filename:
                chapter_id = chapter.get('id', chapter.get('chapter_id', 'unknown'))
                filename = f"chapter_{novel_id}_{chapter_id}.txt"
            
            # 创建小说专属目录
            novel_chapters_dir = self.chapters_dir / novel_id
            novel_chapters_dir.mkdir(parents=True, exist_ok=True)
            
            filepath = novel_chapters_dir / filename
            
            content = f"标题：{chapter.get('title', '')}\n"
            content += f"作者：{chapter.get('author', '')}\n"
            content += f"字数：{chapter.get('word_count', 0)}\n"
            content += f"更新时间：{chapter.get('updated_at', '')}\n"
            content += "=" * 50 + "\n\n"
            content += chapter.get('content', '')
            
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(content)
            
            log_storage_info(f"章节文本保存成功: {filepath}")
            return True
            
        except Exception as e:
            log_crawler_error(f"保存章节文本失败: {str(e)}")
            return False
    
    def save_novel_txt(self, novel: Dict[str, Any], chapters: List[Dict[str, Any]], filename: str = None) -> bool:
        """
        保存完整小说为TXT格式
        
        Args:
            novel: 小说信息
            chapters: 章节列表
            filename: 文件名（可选）
            
        Returns:
            是否保存成功
        """
        try:
            if not filename:
                novel_id = novel.get('id', novel.get('novel_id', 'unknown'))
                filename = f"novel_{novel_id}_complete.txt"
            
            filepath = self.novels_dir / filename
            
            content = f"书名：{novel.get('title', '')}\n"
            content += f"作者：{novel.get('author', '')}\n"
            content += f"分类：{novel.get('category', '')}\n"
            content += f"状态：{novel.get('status', '')}\n"
            content += f"字数：{novel.get('word_count', 0)}\n"
            content += f"评分：{novel.get('rating', 0.0)}\n"
            content += f"简介：{novel.get('summary', '')}\n"
            content += "=" * 80 + "\n\n"
            
            for i, chapter in enumerate(chapters, 1):
                content += f"第{i}章 {chapter.get('title', '')}\n"
                content += "=" * 50 + "\n"
                content += chapter.get('content', '') + "\n\n"
            
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(content)
            
            log_storage_info(f"完整小说TXT保存成功: {filepath} - 共 {len(chapters)} 章")
            return True
            
        except Exception as e:
            log_crawler_error(f"保存完整小说TXT失败: {str(e)}")
            return False
    
    def load_novel_json(self, filename: str) -> Optional[Dict[str, Any]]:
        """
        加载小说JSON数据
        
        Args:
            filename: 文件名
            
        Returns:
            小说数据或None
        """
        try:
            filepath = self.novels_dir / filename
            if not filepath.exists():
                return None
            
            with open(filepath, 'r', encoding='utf-8') as f:
                return json.load(f)
                
        except Exception as e:
            log_crawler_error(f"加载小说JSON失败: {str(e)}")
            return None
    
    def load_chapters_json(self, novel_id: str, filename: str) -> Optional[Dict[str, Any]]:
        """
        加载章节JSON数据
        
        Args:
            novel_id: 小说ID
            filename: 文件名
            
        Returns:
            章节数据或None
        """
        try:
            novel_chapters_dir = self.chapters_dir / novel_id
            filepath = novel_chapters_dir / filename
            
            if not filepath.exists():
                return None
            
            with open(filepath, 'r', encoding='utf-8') as f:
                return json.load(f)
                
        except Exception as e:
            log_crawler_error(f"加载章节JSON失败: {str(e)}")
            return None
    
    def list_novel_files(self, pattern: str = "*.json") -> List[str]:
        """
        列出小说文件
        
        Args:
            pattern: 文件模式
            
        Returns:
            文件列表
        """
        try:
            files = list(self.novels_dir.glob(pattern))
            return [f.name for f in files]
        except Exception as e:
            log_crawler_error(f"列出小说文件失败: {str(e)}")
            return []
    
    def list_chapter_files(self, novel_id: str, pattern: str = "*.json") -> List[str]:
        """
        列出章节文件
        
        Args:
            novel_id: 小说ID
            pattern: 文件模式
            
        Returns:
            文件列表
        """
        try:
            novel_chapters_dir = self.chapters_dir / novel_id
            if not novel_chapters_dir.exists():
                return []
            
            files = list(novel_chapters_dir.glob(pattern))
            return [f.name for f in files]
        except Exception as e:
            log_crawler_error(f"列出章节文件失败: {str(e)}")
            return [] 